{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "709a5ea5-dccb-4128-af56-c58b8c6a4885",
   "metadata": {},
   "source": [
    "### 使用Python\n",
    "\n",
    "\n",
    "#### 训练\n",
    "导入一些库："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "389e31e5-f00c-426c-875e-f85afea9040d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'\n",
    "\n",
    "from swift.llm import get_model_tokenizer, load_dataset, get_template, EncodePreprocessor\n",
    "from swift.utils import get_logger, find_all_linears, get_model_parameter_info, plot_images, seed_everything\n",
    "from swift.tuners import Swift, LoraConfig\n",
    "from swift.trainers import Seq2SeqTrainer, Seq2SeqTrainingArguments\n",
    "from functools import partial\n",
    "\n",
    "logger = get_logger()\n",
    "seed_everything(42)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fec1e09b-3718-40c6-9bc4-83f1abe8ef62",
   "metadata": {},
   "source": [
    "设置训练的超参数："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d986b80b-b4eb-43db-b714-378081a33abb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 模型\n",
    "model_id_or_path = 'Qwen/Qwen2.5-3B-Instruct'  # model_id or model_path\n",
    "system = 'You are a helpful assistant.'\n",
    "output_dir = 'output'\n",
    "\n",
    "# 数据集\n",
    "dataset = ['AI-ModelScope/alpaca-gpt4-data-zh#500', 'AI-ModelScope/alpaca-gpt4-data-en#500',\n",
    "           'swift/self-cognition#500']  # dataset_id or dataset_path\n",
    "data_seed = 42\n",
    "max_length = 2048\n",
    "split_dataset_ratio = 0.01  # 切分验证集\n",
    "num_proc = 4  # 预处理的进程数\n",
    "# 替换自我认知数据集中的填充符：{{NAME}}, {{AUTHOR}}\n",
    "model_name = ['小黄', 'Xiao Huang']  # 模型的中文名和英文名\n",
    "model_author = ['魔搭', 'ModelScope']  # 模型作者的中文名和英文名\n",
    "\n",
    "# lora\n",
    "lora_rank = 8\n",
    "lora_alpha = 32\n",
    "\n",
    "# 训练超参数\n",
    "training_args = Seq2SeqTrainingArguments(\n",
    "    output_dir=output_dir,\n",
    "    learning_rate=1e-4,\n",
    "    per_device_train_batch_size=1,\n",
    "    per_device_eval_batch_size=1,\n",
    "    gradient_checkpointing=True,\n",
    "    weight_decay=0.1,\n",
    "    lr_scheduler_type='cosine',\n",
    "    warmup_ratio=0.05,\n",
    "    report_to=['tensorboard'],\n",
    "    logging_first_step=True,\n",
    "    save_strategy='steps',\n",
    "    save_steps=50,\n",
    "    eval_strategy='steps',\n",
    "    eval_steps=50,\n",
    "    gradient_accumulation_steps=16,\n",
    "    num_train_epochs=1,\n",
    "    metric_for_best_model='loss',\n",
    "    save_total_limit=2,\n",
    "    logging_steps=5,\n",
    "    dataloader_num_workers=1,\n",
    "    data_seed=data_seed,\n",
    ")\n",
    "\n",
    "output_dir = os.path.abspath(os.path.expanduser(output_dir))\n",
    "logger.info(f'output_dir: {output_dir}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba91bf28-fec1-4a80-ba43-10b3f75699c3",
   "metadata": {},
   "source": [
    "获取模型和对话template，并将可训练的lora层加入到模型中："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ade53c46-aef5-4cde-83c4-1d7ba781a155",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model, tokenizer = get_model_tokenizer(model_id_or_path)\n",
    "logger.info(f'model_info: {model.model_info}')\n",
    "template = get_template(model.model_meta.template, tokenizer, default_system=system, max_length=max_length)\n",
    "template.set_mode('train')\n",
    "\n",
    "target_modules = find_all_linears(model)\n",
    "lora_config = LoraConfig(task_type='CAUSAL_LM', r=lora_rank, lora_alpha=lora_alpha,\n",
    "                         target_modules=target_modules)\n",
    "model = Swift.prepare_model(model, lora_config)\n",
    "logger.info(f'lora_config: {lora_config}')\n",
    "\n",
    "# 打印模型结构和训练的参数量\n",
    "logger.info(f'model: {model}')\n",
    "model_parameter_info = get_model_parameter_info(model)\n",
    "logger.info(f'model_parameter_info: {model_parameter_info}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca2b77fe-5048-4323-9a57-c1f7a4042334",
   "metadata": {},
   "source": [
    "下载并载入数据集，并切分成训练集和验证集，\n",
    "\n",
    "然后将文本编码成tokens："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "638d3513-e734-4307-b070-178a4f0128b6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "train_dataset, val_dataset = load_dataset(dataset, split_dataset_ratio=split_dataset_ratio, num_proc=num_proc,\n",
    "        model_name=model_name, model_author=model_author, seed=data_seed)\n",
    "\n",
    "logger.info(f'train_dataset: {train_dataset}')\n",
    "logger.info(f'val_dataset: {val_dataset}')\n",
    "logger.info(f'train_dataset[0]: {train_dataset[0]}')\n",
    "\n",
    "train_dataset = EncodePreprocessor(template=template)(train_dataset, num_proc=num_proc)\n",
    "val_dataset = EncodePreprocessor(template=template)(val_dataset, num_proc=num_proc)\n",
    "logger.info(f'encoded_train_dataset[0]: {train_dataset[0]}')\n",
    "\n",
    "# 打印一条样本\n",
    "template.print_inputs(train_dataset[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4f90f0e-8da2-42ab-8f5b-2a39da539a5e",
   "metadata": {},
   "source": [
    "初始化trainer并开始训练："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "959b1a9c-cf0b-4985-9dac-1bc9b1d12473",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model.enable_input_require_grads()  # 兼容gradient checkpointing\n",
    "trainer = Seq2SeqTrainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    data_collator=template.data_collator,\n",
    "    train_dataset=train_dataset,\n",
    "    eval_dataset=val_dataset,\n",
    "    template=template,\n",
    ")\n",
    "trainer.train()\n",
    "\n",
    "last_model_checkpoint = trainer.state.last_model_checkpoint\n",
    "logger.info(f'last_model_checkpoint: {last_model_checkpoint}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d178b52-fedb-4819-a463-27b624f310ff",
   "metadata": {},
   "source": [
    "可视化训练的loss。其中浅黄色线条代表真实loss值，黄色线条代表经过0.9平滑系数平滑后的loss值。\n",
    "\n",
    "你也可以使用tensorboard进行实时可视化，在命令行输入`tensorboard --logdir '{output_dir}/runs'`。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92274822-0837-4521-9ad9-840ec0425d85",
   "metadata": {},
   "outputs": [],
   "source": [
    "images_dir = os.path.join(output_dir, 'images')\n",
    "logger.info(f'images_dir: {images_dir}')\n",
    "plot_images(images_dir, training_args.logging_dir, ['train/loss'], 0.9)  # 保存图片\n",
    "\n",
    "# 展示图片\n",
    "from IPython.display import display\n",
    "from PIL import Image\n",
    "image = Image.open(os.path.join(images_dir, 'train_loss.png'))\n",
    "display(image)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80b2216b-02a7-45f2-88ec-0d5c7b273de9",
   "metadata": {},
   "source": [
    "#### 微调后推理"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdb6217b-2647-41f4-b6cc-32275dd829f3",
   "metadata": {},
   "source": [
    "导入一些库："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32208bd0-806e-4de5-86fc-1017366ab4f8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'\n",
    "\n",
    "from swift.llm import InferEngine, InferRequest, PtEngine, RequestConfig, get_template"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "043c0967-2a57-4dc4-aaa0-62e23f53d3cd",
   "metadata": {},
   "source": [
    "设置推理的超参数："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0229fa94-1aab-4727-930d-e354bdf30986",
   "metadata": {
    "ExecutionIndicator": {
     "show": false
    },
    "execution": {
     "iopub.execute_input": "2024-12-25T16:56:32.156242Z",
     "iopub.status.busy": "2024-12-25T16:56:32.155674Z",
     "iopub.status.idle": "2024-12-25T16:56:32.159195Z",
     "shell.execute_reply": "2024-12-25T16:56:32.158698Z",
     "shell.execute_reply.started": "2024-12-25T16:56:32.156217Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "last_model_checkpoint = 'output/vx-xxx/checkpoint-xxx'\n",
    "\n",
    "# 模型\n",
    "model_id_or_path = 'Qwen/Qwen2.5-3B-Instruct'  # model_id or model_path\n",
    "system = 'You are a helpful assistant.'\n",
    "infer_backend = 'pt'\n",
    "\n",
    "# 生成参数\n",
    "max_new_tokens = 512\n",
    "temperature = 0\n",
    "stream = True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da8a3bf9-c321-4f33-9251-944059537a4d",
   "metadata": {},
   "source": [
    "获取推理引擎，并载入LoRA权重："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47bcc49e-c354-48bf-b869-3db603e8d648",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "engine = PtEngine(model_id_or_path, adapters=[last_model_checkpoint])\n",
    "template = get_template(engine.model.model_meta.template, engine.tokenizer, default_system=system)\n",
    "# 这里对推理引擎的默认template进行修改，也可以在`engine.infer`时进行传入\n",
    "engine.default_template = template"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71a01fa0-673d-44f3-9891-b31faaf63b6f",
   "metadata": {},
   "source": [
    "开始推理..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d657572-67a1-4224-ac37-6fefee54d179",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "query_list = [\n",
    "    'who are you?',\n",
    "    \"晚上睡不着觉怎么办？\",\n",
    "    '你是谁训练的？',\n",
    "]\n",
    "\n",
    "def infer_stream(engine: InferEngine, infer_request: InferRequest):\n",
    "    request_config = RequestConfig(max_tokens=max_new_tokens, temperature=temperature, stream=True)\n",
    "    gen_list = engine.infer([infer_request], request_config)\n",
    "    query = infer_request.messages[0]['content']\n",
    "    print(f'query: {query}\\nresponse: ', end='')\n",
    "    for resp in gen_list[0]:\n",
    "        if resp is None:\n",
    "            continue\n",
    "        print(resp_list[0].choices[0].delta.content, end='', flush=True)\n",
    "    print()\n",
    "\n",
    "def infer(engine: InferEngine, infer_request: InferRequest):\n",
    "    request_config = RequestConfig(max_tokens=max_new_tokens, temperature=temperature)\n",
    "    resp_list = engine.infer([infer_request], request_config)\n",
    "    query = infer_request.messages[0]['content']\n",
    "    response = resp_list[0].choices[0].message.content\n",
    "    print(f'query: {query}')\n",
    "    print(f'response: {response}')\n",
    "\n",
    "infer_func = infer_stream if stream else infer\n",
    "for query in query_list:\n",
    "    infer_func(engine, InferRequest(messages=[{'role': 'user', 'content': query}]))\n",
    "    print('-' * 50)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rs_bigmodel",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
